#检索器
"""
langchain中使用chroma
1.构建Document list
2.构建embeddings模型对象
3.构建向量数据库 Chroma.from_documents(docs,emb,persist_directory="chromadb")  langchain_core.documents import Document
4.构建检索器 vector_db.as_retriever
5.执行检索器查询  retriever.invoke("马云的花名是什么")
"""

from model_utils import getEmd
from langchain_chroma import Chroma
from langchain_core.documents import Document

text = """
基本信息
姓名：马云（英文名：Jack Ma）
出生日期：1964年9月10日
出生地：中国浙江省杭州市
毕业院校：杭州师范大学（原杭州师范学院），英语专业
花名：风清扬（源自金庸小说《笑傲江湖》中的神秘高手风清扬）
"""

vector_db = Chroma.from_documents([Document(text)], getEmd(), persist_directory="./chromadb_1")

#retriever = vector_db.as_retriever(search_kwargs={"k":5}) #返回最相关的5条文档
retriever = vector_db.as_retriever(search_type="similarity_score_threshold",search_kwargs={'score_threshold': 0.5}) #只返回相似度得分超过指定阈值（0.5）的文档

r = retriever.invoke("马云的花名是什么")
print(r)
print(len(r))